import { useCallback, useRef, useState } from "react";
import { API_BASE_URL } from "../constants";
import { request } from "shared/requests";
import { ReadableStreamParser, JsonParser } from "shared/helpers";

class CancelRequestError extends Error {
  name: "CancelRequest";
  message: "Cancel";
}

type CompletionData = {
  uri: string;
  model: string;
  year: string;
  question?: string;
};

export default function useCompletion() {
  const [isLoading, setIsLoading] = useState(false);
  const [error, setError] = useState<Error>(null);
  const [chunks, setChunks] = useState([]);
  const [params, setParams] = useState<CompletionData>({
    uri: "",
    model: "",
    year: "",
    question: undefined,
  });
  const abortControllerRef = useRef<AbortController>();
  const handleCompletion = useCallback((data: CompletionData) => {
    const { uri, model, year, question } = data;
    if (!uri || !model || !year) {
      return;
    }
    setParams(data);
    if (abortControllerRef.current) {
      abortControllerRef.current.abort(new CancelRequestError());
    }
    abortControllerRef.current = new AbortController();
    setError(null);
    setIsLoading(true);
    setChunks((chunks) => (chunks.length === 0 ? chunks : []));
    completion({
      data: { uri, model, year, question },
      signal: abortControllerRef.current.signal,
      onProgress: (chunk) => {
        setChunks((chunks) => [...chunks, chunk]);
      },
    })
      .then(() => {
        setIsLoading(false);
      })
      .catch((err) => {
        if (
          !(err instanceof CancelRequestError) &&
          err.message !== "BodyStreamBuffer was aborted"
        ) {
          setError(err);
          setIsLoading(false);
        }
      });
  }, []);

  const retry = () => handleCompletion(params);

  return { isLoading, error, chunks, handleCompletion, retry, params };
}

async function completion(options: {
  data: { uri: string; model: string; year: string; question?: string };
  onProgress?: (v: any) => void;
  signal: AbortSignal;
}) {
  const { data, onProgress, signal } = options;

  return request(`${API_BASE_URL}/query`, {
    method: "post",
    headers: { "content-type": "application/json" },
    data,
    signal,
    stream: true,
  }).then(async (stream) => {
    const parser = new ReadableStreamParser();
    parser.use(
      new JsonParser({
        onParsed: onProgress,
      }),
    );
    parser.parse(stream);
    return new Promise((resolve, reject) => {
      parser.on("end", resolve);
      parser.on("error", reject);
    });
  });
}
